
library(doMC); registerDoMC()
library(random)
library(doRNG)

library(endorse)

rm(list=ls())

load("dataSmall-violence-clean.RData")
load("extras.RData")

set.seed(685994)

drop <- "Q1" ## change to "Q2", "Q3", or "Q4" to drop the other policies

qlist <- c("Q1", "Q2", "Q3", "Q4")
qlist <- qlist[-which(qlist==drop)]

treat.ISAFControl <- endorser[,1] == 1 | endorser[,1] == 2

dat <- as.data.frame(cbind(Y, indivcov, data[, c(colnames(villcov), colnames(distcov))]))[treat.ISAFControl == TRUE, ]

names(dat)[1:4] <- paste("Q", 1:4, sep="")

formula.temp <- as.formula(paste(" ~ ", paste(paste(colnames(indivcov), collapse = " + "), " + ", paste(colnames(villcov), collapse = " + "), " + ", paste(colnames(distcov), collapse = " + "))))

cov.mat.temp <- model.matrix(formula.temp, dat)

theta.start <- delta.start <- list()
lambda.start <- list()
x.start <- list()
for(i in 1:3) {
  theta.start[[i]] <- matrix(rnorm(3*29, mean = 0, sd = 5), nrow = 3, ncol = 29)
  theta.temp <- theta.start[[i]]
  lambda.start[[i]] <- matrix(rnorm(1*29, mean = 0, sd = 5), nrow = 1, ncol = 29)
  omega.start <- 1
  phi2.start <- 1
  delta.start[[i]] <- rnorm(29, mean = 0, sd = 5)
  x.start[[i]] <- rnorm(nrow(dat),
                        mean = as.numeric(cov.mat.temp %*% delta.start[[i]]),
                        sd = sqrt(9))
}

endorser <- endorser[treat.ISAFControl == TRUE,] - 1

MCMC <- 60000
burn <- 30000
thin <- 100

procs <- 3

set.seed(45622)

res <- foreach(i=1:procs) %dorng% {
  result <- endorse(Y = list(Q1 = qlist[1], Q2 = qlist[2], Q3 = qlist[3]),
                      data = dat, treat = endorser, covariates = TRUE,
                      formula.indiv = formula.temp,
                      delta.start = delta.start[[i]],
                      x.start = x.start[[i]], lambda.start = lambda.start[[i]],
                      precision.beta = 1/9, precision.x = 1, precision.theta = 1/9,
                      precision.delta = 1/9,
                      tau.out = TRUE, MCMC = MCMC, burn = burn, thin = thin,
                      s.out = TRUE, x.sd = FALSE, 
                      seed.store = TRUE, identical.lambda = TRUE)
 
  return(result)
}

save(res, file = paste("../data/results-endorse-villdist-AJPS-parallel-drop", drop, ".RData", sep = ""))

J <- 3
M <- 29
K <- 1

intercept1 <- matrix(NA, nrow = floor((MCMC - burn) / thin), ncol = 3)
colnames(intercept1) <- paste("intercept", 1:3, sep = ".")
for (j in 1:3) {
  intercept1[, j] <- - res[[1]]$beta[, (2 * j - 1)] +
    res[[1]]$delta[, 1] * res[[1]]$beta[, (2 * j)]
}
intercept2 <- matrix(NA, nrow = floor((MCMC - burn) / thin), ncol = 3)
colnames(intercept2) <- paste("intercept", 1:3, sep = ".")
for (j in 1:3) {
  intercept2[, j] <- - res[[2]]$beta[, (2 * j - 1)] +
    res[[2]]$delta[, 1] * res[[2]]$beta[, (2 * j)]
}
intercept3 <- matrix(NA, nrow = floor((MCMC - burn) / thin), ncol = 3)
colnames(intercept3) <- paste("intercept", 1:3, sep = ".")
for (j in 1:3) {
  intercept3[, j] <- - res[[3]]$beta[, (2 * j - 1)] +
    res[[3]]$delta[, 1] * res[[3]]$beta[, (2 * j)]
}

chain1 <- mcmc(cbind(intercept1, res[[1]]$beta[, (1:J) * 2],
                     res[[1]]$lambda, 
                     res[[1]]$delta[, 2:M], res[[1]]$omega2,
                     res[[1]]$phi2), start = burn + 1, thin = thin)
chain2 <- mcmc(cbind(intercept2, res[[2]]$beta[, (1:J) * 2],
                     res[[2]]$lambda, 
                     res[[2]]$delta[, 2:M], res[[2]]$omega2,
                     res[[2]]$phi2), start = burn + 1, thin = thin)
chain3 <- mcmc(cbind(intercept3, res[[3]]$beta[, (1:J) * 2],
                     res[[3]]$lambda, 
                     res[[3]]$delta[, 2:M], res[[3]]$omega2,
                     res[[3]]$phi2), start = burn + 1, thin = thin)

check.R.hat <- mcmc.list(chain1, chain2, chain3)

R.hat <- gelman.diag(check.R.hat)

while (max(R.hat$psrf[, 2]) > 1.1) {
  burn <- 0

  res.last <- res
  
  res <- foreach(i=1:procs) %dorng% {
      result <- endorse(Y = list(Q1 = qlist[1], Q2 = qlist[2], Q3 = qlist[3]),
                      data = dat, treat = endorser, covariates = TRUE,
                      formula.indiv = formula.temp,
                      precision.beta = 1/9, precision.x = 1, precision.theta = 1/9,
                      precision.delta = 1/9,
                      tau.out = TRUE, s.out = TRUE, x.sd = FALSE, MCMC = MCMC, burn = burn, thin = thin,
                      seed.store = TRUE, identical.lambda = TRUE,
                      update = TRUE, update.start = res.last[[i]])
    return(result)
  }

  for(i in 1:procs) {
    res[[i]]$x <- res[[i]]$x[nrow(res[[i]]$x), , drop = FALSE]
    res[[i]]$s <- res[[i]]$s[nrow(res[[i]]$s), , drop = FALSE]
    res[[i]]$beta <- rbind(res.last[[i]]$beta, res[[i]]$beta)
    res[[i]]$tau <- rbind(res.last[[i]]$tau, res[[i]]$tau)
    res[[i]]$lambda <- rbind(res.last[[i]]$lambda, res[[i]]$lambda)
    res[[i]]$omega2 <- rbind(res.last[[i]]$omega2, res[[i]]$omega2)
    res[[i]]$delta <- rbind(res.last[[i]]$delta, res[[i]]$delta)
    res[[i]]$phi2 <- rbind(res.last[[i]]$phi2, res[[i]]$phi2)
  }

  rows <- nrow(res[[i]]$lambda)
  
  intercept1 <- matrix(NA, nrow = rows, ncol = 3)
  colnames(intercept1) <- paste("intercept", 1:3, sep = ".")
  for (j in 1:3) {
    intercept1[, j] <- - res[[1]]$beta[, (2 * j - 1)] +
      res[[1]]$delta[, 1] * res[[1]]$beta[, (2 * j)]
  }
  intercept2 <- matrix(NA, nrow = rows, ncol = 3)
  colnames(intercept2) <- paste("intercept", 1:3, sep = ".")
  for (j in 1:3) {
    intercept2[, j] <- - res[[2]]$beta[, (2 * j - 1)] +
      res[[2]]$delta[, 1] * res[[2]]$beta[, (2 * j)]
  }
  intercept3 <- matrix(NA, nrow = rows, ncol = 3)
  colnames(intercept3) <- paste("intercept", 1:3, sep = ".")
  for (j in 1:3) {
    intercept3[, j] <- - res[[3]]$beta[, (2 * j - 1)] +
      res[[3]]$delta[, 1] * res[[3]]$beta[, (2 * j)]
  }

  chain1 <- mcmc(cbind(intercept1, res[[1]]$beta[, (1:J) * 2],
                       res[[1]]$lambda, 
                       res[[1]]$delta[, 2:M], res[[1]]$omega2,
                       res[[1]]$phi2), start = burn + 1, thin = thin)
  chain2 <- mcmc(cbind(intercept2, res[[2]]$beta[, (1:J) * 2],
                       res[[2]]$lambda,
                       res[[2]]$delta[, 2:M], res[[2]]$omega2,
                       res[[2]]$phi2), start = burn + 1, thin = thin)
  chain3 <- mcmc(cbind(intercept3, res[[3]]$beta[, (1:J) * 2],
                       res[[3]]$lambda, 
                       res[[3]]$delta[, 2:M], res[[3]]$omega2,
                       res[[3]]$phi2), start = burn + 1, thin = thin)

  check.R.hat <- mcmc.list(chain1, chain2, chain3)

  R.hat <- gelman.diag(check.R.hat)

  R.hat

  save(res, file = paste("../data/results-endorse-villdist-AJPS-parallel-drop", drop, ".RData", sep = ""))
  
}

